# models/train_eval.py
import paddle
import paddle.nn as nn
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


def train_model(model, X, time_X, y, epochs=100, batch_size=32, curriculum_steps=None):
    """
    Curriculum Learning 版本训练函数

    参数:
    - model: 模型
    - X: 历史序列 [N, seq_len, n_vars]
    - time_X: 时间特征 [N, seq_len, 2]
    - y: 预测目标 [N, pred_len]
    - epochs: 总训练轮数
    - batch_size: 批量大小
    - curriculum_steps: list of (pred_len, top_k)，如 [(6, 2), (12, 3), (24, 5)]
    """
    if curriculum_steps is None:
        curriculum_steps = [(6, 2), (12, 3), (24, 5)]

    X_tensor = paddle.to_tensor(X, dtype='float32')
    time_tensor = paddle.to_tensor(time_X, dtype='int64')
    y_tensor = paddle.to_tensor(y, dtype='float32')

    # 初始化优化器
    optimizer = paddle.optimizer.AdamW(
        parameters=model.parameters(),
        learning_rate=1e-3,
        weight_decay=1e-5
    )
    scheduler = paddle.optimizer.lr.ReduceOnPlateau(
        learning_rate=optimizer.get_lr(),
        mode='min',
        factor=0.5,
        patience=5
    )

    best_loss = float('inf')
    train_losses = []

    for pred_len, top_k in curriculum_steps:
        print(f"\n[Curriculum Learning] 当前训练目标: {pred_len} 步，top_k={top_k}")
        # 动态修改预测头长度
        d_model = model.head[0].weight.shape[0]  # 获取 d_model 维度
        model.head = nn.Sequential(
            nn.Linear(d_model, d_model * 2),
            nn.GELU(),
            nn.Linear(d_model * 2, pred_len)
        )

        # 重新构建数据集（因为 pred_len 变了）
        y_tensor_curriculum = y_tensor[:, :pred_len]

        for epoch in range(epochs // len(curriculum_steps)):
            model.train()
            epoch_loss = 0
            permutation = np.random.permutation(len(X))

            for i in range(0, len(X), batch_size):
                indices = permutation[i:i + batch_size]
                batch_X = X_tensor[indices]
                batch_time = time_tensor[indices]
                batch_y = y_tensor_curriculum[indices]

                pred = model(batch_X, batch_time)

                # ✅ 取 pred 的最后一个时间步，使其与 batch_y 形状一致
                loss = nn.functional.mse_loss(pred[:, -1, :], batch_y)

                loss.backward()
                optimizer.step()
                optimizer.clear_grad()

                epoch_loss += loss.numpy()

            avg_loss = epoch_loss / (len(X) / batch_size)
            train_losses.append(avg_loss)
            scheduler.step(avg_loss)

            if avg_loss < best_loss:
                best_loss = avg_loss
                paddle.save(model.state_dict(), 'best_model.pdparams')
                print(f"Saved best model with loss: {best_loss:.6f}")

            print(f"Epoch {epoch + 1}, Pred_len={pred_len}, Loss={avg_loss:.6f}")
            print("="*13)

    # 绘图保存
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label='Training Loss')
    plt.title('Training Loss Over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('MSE Loss')
    plt.legend()
    plt.savefig('training_loss.png')
    plt.close()

    return train_losses


def evaluate_model(model, X_test, time_X_test, y_test, scaler, return_predictions=False):
    """
    模型评估函数

    参数:
    - model: 模型
    - X_test: 测试集历史数据 [N, seq_len, n_vars]
    - time_X_test: 时间特征 [N, seq_len, 2]
    - y_test: 真实值 [N, pred_len]
    - scaler: 数据标准化器
    - return_predictions: 是否返回预测值用于后续分析

    返回:
    - test_loss: 测试损失
    - y_pred: 预测值（如果 return_predictions=True）
    """
    model.eval()
    with paddle.no_grad():
        X_tensor = paddle.to_tensor(X_test, dtype='float32')
        time_tensor = paddle.to_tensor(time_X_test, dtype='int64')
        y_tensor = paddle.to_tensor(y_test, dtype='float32')

        pred = model(X_tensor, time_tensor)

        # ✅ 取 pred 的最后一个时间步进行评估
        loss = nn.functional.mse_loss(pred[:, -1, :], y_tensor).numpy()

        # 反归一化
        y_true = scaler.inverse_transform(y_test[0].reshape(-1, 1)).flatten()
        y_pred = scaler.inverse_transform(pred[0].numpy().reshape(-1, 1)).flatten()

        # 可视化
        plt.figure(figsize=(15, 6))
        plt.plot(y_true, label='True')
        plt.plot(y_pred, label='Predicted')
        plt.title(f"Time Series Prediction (MSE: {loss:.4f})")
        plt.xlabel('Time Steps')
        plt.ylabel('CPU Usage')
        plt.legend()
        plt.grid(True)
        plt.savefig('prediction_result.png')
        plt.close()

        if return_predictions:
            return loss, y_pred
        return loss
